import cv2
import os
import json
import numpy as np
import time # 用于测试性能

BASE_DIR = os.path.dirname(os.path.abspath(__file__))
CONFIG_FILE = os.path.join(BASE_DIR, "..", "model", "thresholds.json")
config_path = os.path.abspath(CONFIG_FILE)

def build_hsv_map(config_list):
    hsv_map = {}
    for cfg in config_list:
        mode = cfg['mode']
        name = cfg['color']
        params = cfg['params']
        if mode == 'hsv':
            hsv_map[name] = {
                'lower': (params['h_low'], params['s_low'], params['v_low']),
                'upper': (params['h_high'], params['s_high'], params['v_high'])
            }
        elif mode == 'binary':
            hsv_map[name] = {
                'thresh': params['thresh']
            }
        else:
            raise ValueError(f"不支持的模式: {mode}")
    return hsv_map

try:
    with open(config_path, 'r') as f:
        hsv_config = json.load(f)
except FileNotFoundError:
    print(f"Error: Config file not found at {config_path}")
    hsv_config = {} # Handle error appropriately

hsv_config = build_hsv_map(hsv_config)




# def extract_white_roi_perspective(img, min_area=2000, dst_size=(640, 480)):
#     """
#     找到图中最大的白色区域，对该轮廓做透视变换，返回一张 dst_size 大小的 BGR 图。
#     如果 dst_size=None，则用该轮廓的 boundingRect 尺寸。
#     """
#     global last_points, prev_src
#     # 1. 提取白色二值 mask
#     hsv   = cv2.cvtColor(img, cv2.COLOR_BGR2HSV)
#     lower = np.array(hsv_config['white']['lower'])
#     upper = np.array(hsv_config['white']['upper'])
#     mask  = cv2.inRange(hsv, lower, upper)
#     mask =cv2.medianBlur(mask,5)

#     # 2. 查找轮廓并过滤小面积
#     cnts, _ = cv2.findContours(mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
#     cnts    = [c for c in cnts if cv2.contourArea(c) > min_area]
#     if not cnts:
#         return None  # 没有找到足够大的白色区域

#     # 3. 选取最大轮廓
#     c = max(cnts, key=cv2.contourArea)

#     # 4. 拟合四边形顶点
#     peri   = cv2.arcLength(c, True)
#     approx = cv2.approxPolyDP(c, 0.02 * peri, True)
#     if len(approx) == 4:
#         pts = approx.reshape(4, 2).astype(np.float32)
#     else:
#         # fallback：使用最小外接矩形
#         x, y, w, h = cv2.boundingRect(c)
#         pts = np.array([[x,y], [x+w,y], [x+w,y+h], [x,y+h]], np.float32)

#     # 5. 目标输出大小
#     if dst_size is None:
    
#         x, y, w, h = cv2.boundingRect(c)
#         dst_w, dst_h = w, h
#     else:
#         dst_w, dst_h = dst_size

#     # 6. 点排序函数：保证顺序为 tl, tr, br, bl
#     def order_points(pts):
#         s    = pts.sum(axis=1)
#         diff = np.diff(pts, axis=1)
#         tl   = pts[np.argmin(s)]
#         br   = pts[np.argmax(s)]
#         tr   = pts[np.argmin(diff)]
#         bl   = pts[np.argmax(diff)]
#         return np.array([tl, tr, br, bl], np.float32)
    
#     src = order_points(pts)
#     # 1. 如果有历史点，则做平滑
#     if len(last_points) > 2:
#         avg_src = np.mean([last_points[-1], last_points[-2], last_points[-3]], axis=0)
#         src= src * 0.7 + avg_src * 0.3
#      # 2. 判断与上次角点的最大偏移
#     if prev_src is not None:
#         # 计算4个角点的最大距离
#         max_dist = np.max(np.linalg.norm(src - prev_src, axis=1))
#         threshold = 5.0  # 你可以调节这个阈值
#         if max_dist < threshold:
#             # 偏移小，认为角点没变化，复用上次
#             src = prev_src.copy()
#         else:
#             # 偏移大，更新prev_src
#             prev_src = src.copy()
#     else:
#         prev_src = src.copy()

#     last_points.append(src)

#     if len(last_points) > 3:
#         last_points.pop(0)

#     dst = np.array([[0, 0], [dst_w-1, 0], [dst_w-1, dst_h-1], [0, dst_h-1]], np.float32)
#     # 7. 计算透视矩阵并变换
#     M      = cv2.getPerspectiveTransform(src, dst)
#     warped = cv2.warpPerspective(img, M, (dst_w, dst_h))

#     return warped
def extract_white_roi_perspective(img, min_area=20000):
    """
    提取非黑色区域的ROI（以白色为主），并返回ROI和其包围盒。
    修复掩膜类型问题，确保与bitwise_and兼容。
    """
    # 转为灰度图
    gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
    thresh_low = np.array(hsv_config['thresh_white']['thresh'])

    # 对图像进行二值化处理，阈值可调：靠近255为更亮区域
    _, binary = cv2.threshold(gray, int(thresh_low), 255, cv2.THRESH_BINARY)

    # 寻找轮廓
    contours, _ = cv2.findContours(binary, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
    valid_contours = [cnt for cnt in contours if cv2.contourArea(cnt) > min_area]

    if not valid_contours:
        return None, None, None

    # 取最大轮廓
    largest_contour = max(valid_contours, key=cv2.contourArea)

    # 创建掩膜，类型需为uint8且单通道
    mask = np.zeros(img.shape[:2], dtype=np.uint8)
    cv2.drawContours(mask, [largest_contour], -1, (255,), thickness=cv2.FILLED)
    # 保存调试图
    cv2.imwrite("mask.jpg", mask)
    # 用掩膜提取原图区域
    roi = cv2.bitwise_and(img, img, mask=mask)
    # 裁剪 mask 覆盖区域的最小包围矩形
    x, y, w, h = cv2.boundingRect(largest_contour)
    cropped_roi = roi[y:y+h, x:x+w]  # 用于显示/保存

    return roi


def preprocess_mask(hsv, lower, upper, kernel_size=(5,5), open_iterations=2, blur_kernel=(5,5)):
    """
    预处理掩码

    Args:
        hsv: HSV 图像
        lower: HSV 范围下限
        upper: HSV 范围上限
        kernel_size: 形态学操作核大小
        open_iterations: 开运算迭代次数
        blur_kernel: 高斯模糊核大小

    Returns:
        处理后的掩码
    """
    mask = cv2.inRange(hsv, lower, upper)
    kernel = np.ones(kernel_size, np.uint8) # 在函数内创建 kernel，更灵活
    mask = cv2.dilate(mask, kernel, iterations=open_iterations)
    # mask = cv2.morphologyEx(mask, cv2.MORPH_OPEN, kernel, iterations=open_iterations)
    if blur_kernel and blur_kernel[0] > 0 and blur_kernel[1] > 0:
         mask = cv2.GaussianBlur(mask, blur_kernel, 0)
    return mask

def find_color_contour(mask, min_area=1000,max_area=10000):
    """
    在掩码中查找最大的轮廓

    Args:
        mask: 二值掩码
        min_area: 最小轮廓面积阈值
        max_area: 最大轮廓面积阈值

    Returns:
        (轮廓, (cx, cy)) 或 (None, None)
    """
    # cv2.findContours 在 OpenCV 4+ 返回两个值，在 OpenCV 3 返回三个值
   
    contours, _ = cv2.findContours(mask,
                                   cv2.RETR_EXTERNAL, # 只找外部轮廓，更快
                                   cv2.CHAIN_APPROX_SIMPLE) # 压缩点，更快

    if not contours:
        return None, None
    # 找到面积最大的轮廓
    cnt = max(contours, key=cv2.contourArea)
    area = cv2.contourArea(cnt)
    # print(f"Found contour with area {area}")

    if area < min_area or area > max_area:
        return None, None

    # 计算轮廓的中心
    M = cv2.moments(cnt)
    if M['m00'] == 0: # 防止除以零
        return None, None

    cx = int(M['m10']/M['m00'])
    cy = int(M['m01']/M['m00'])

    return cnt, (cx, cy)


def hsv_find(img, target_color=None, min_area=1000, max_area=10000, 
                       scale_factor=1.0, # 新增：缩放因子 (例如 0.5 表示缩小一半)
                       hsv_config=hsv_config, # 预加载的配置
                       kernel_size=(5,5), open_iterations=2, blur_kernel=(5,5) # 形态学/模糊参数
                       ):
    """
    使用优化的方法在图像中查找指定颜色的中心。

    img: 输入 BGR 图像
    target_color: None 或 'red'、'blue'、'green'、'yellow'
    scale_factor: 图像缩放因子。1.0 表示不缩放，<1.0 表示缩小，>1.0 表示放大。
    hsv_config: 预加载的 HSV 配置字典。
    kernel_size: 形态学操作核大小。
    open_iterations: 开运算迭代次数。
    blur_kernel: 高斯模糊核大小 (None 或 (0,0) 表示不模糊)。

    返回： (颜色名称, (cx, cy))，或 (None, None) 如果没有找到
    """
   
    if target_color not in hsv_config:
        print(f"Error: Configuration for color '{target_color}' not found.")
        return None, None, None
    # 1. 可选：图像缩放
    if scale_factor != 1.0:
        # 使用 cv2.INTER_AREA 进行缩小通常效果更好
        img_processed = cv2.resize(img, None, fx=scale_factor, fy=scale_factor, interpolation=cv2.INTER_AREA if scale_factor < 1.0 else cv2.INTER_LINEAR)
    else:
        img_processed = img

    # 2. BGR 转 HSV
    hsv = cv2.cvtColor(img_processed, cv2.COLOR_BGR2HSV)
    mask = None

    if target_color == 'red':
        lower_1 = np.array(hsv_config['redtwo']['lower'])
        upper_1 = np.array(hsv_config['redtwo']['upper'])
        # 3. 获取 HSV 阈值 (使用预加载的配置)
        lower_2 = np.array(hsv_config[target_color]['lower'])
        upper_2 = np.array(hsv_config[target_color]['upper'])
        mask_1 = preprocess_mask(hsv, lower_1, upper_1, kernel_size=kernel_size, open_iterations=open_iterations, blur_kernel=blur_kernel)
        mask_2 = preprocess_mask(hsv, lower_2, upper_2, kernel_size=kernel_size, open_iterations=open_iterations, blur_kernel=blur_kernel)
        mask = cv2.bitwise_or(mask_1, mask_2)
    else:
        # 3. 获取 HSV 阈值 (使用预加载的配置)
        lower = np.array(hsv_config[target_color]['lower'])
        upper = np.array(hsv_config[target_color]['upper'])
        mask = preprocess_mask(hsv, lower, upper, kernel_size=kernel_size,
                           open_iterations=open_iterations, blur_kernel=blur_kernel)

    # 5. 查找颜色轮廓
    # 注意：面积阈值也应该按比例缩放
    scaled_min_area_thresh = int(min_area * (scale_factor ** 2))
    scaled_max_area_thresh = int(max_area * (scale_factor ** 2))
    cnt, center = find_color_contour(mask, scaled_min_area_thresh, scaled_max_area_thresh)

    if cnt is not None and center is not None:
        cx, cy = center
        # 6. 将中心坐标按比例还原回原始图像尺寸
        if scale_factor != 1.0:
            cx = int(cx / scale_factor)
            cy = int(cy / scale_factor)

        # 返回原始图像 (或处理过的图像), 掩码, 和还原后的中心坐标
        return img, mask, (cx, cy) # 或者返回 img_processed, mask, (cx, cy) 如果需要处理过的图

    # 没找到任何颜色
    return None, None, None


# if __name__ == '__main__':
#     # 测试
#     img_path = '/home/sunrise/train1/images/3.jpg'
#     img = cv2.imread(img_path)
#     start_time = time.time()
#     img, mask, center = hsv_find(img, target_color='blue', min_area=100, max_area=10000)
#     end_time = time.time()
#     print(f"Found center: {center} in {end_time - start_time:.3f} seconds")
    
#     if center is not None and img is not None:
#         cv2.circle(img, center, 5, (0, 0, 255), -1)
    
#     if mask is not None:
#         cv2.imshow('mask', mask)
#     if img is not None:
#         cv2.imshow('img', img)
#     cv2.waitKey(0)
#     cv2.destroyAllWindows()
